package com.gitee.wsl.flow.combine


import com.gitee.wsl.ext.any.NULL_VALUE
import com.gitee.wsl.ext.coroutines.set
import kotlinx.atomicfu.atomic
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CoroutineStart
import kotlinx.coroutines.cancel
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.flow.Flow
import kotlinx.coroutines.flow.flow
import kotlinx.coroutines.launch


/**
 * Merges two [Flow]s into one [Flow] by combining each value from self with the latest value from the second [Flow], if any.
 * Values emitted by self before the second [Flow] has emitted any values will be omitted.
 *
 * @param other Second [Flow]
 * @param transform A transform function to apply to each value from self combined with the latest value from the second [Flow], if any.
 */
fun <A, B, R> Flow<A>.withLatestFrom(
    other: Flow<B>,
    transform: suspend (A, B) -> R
): Flow<R> {
    return flow {
        val otherRef = atomic<Any?>(null)

        try {
            coroutineScope {
                launch(start = CoroutineStart.UNDISPATCHED) {
                    other.collect { otherRef.value = it ?: NULL_VALUE }
                }

                collect { value ->
                    emit(
                        transform(
                            value,
                            NULL_VALUE.unbox(otherRef.value ?: return@collect)
                        )
                    )
                }
            }
        } finally {
            otherRef.value = null
        }
    }
}

/*
fun <A, B : Any?, R> Flow<A>.withLatestFrom(
    other: Flow<B>,
    transform: suspend (A, B) -> R
): Flow<R> = flow {
    coroutineScope {
        val latestB = atomic<B?>(null)
        val outerScope = this
        launch {
            try {
                other.collect { latestB.set(it) }
            } catch (e: CancellationException) {
                outerScope.cancel(e) // cancel outer scope on cancellation exception, too
            }
        }
        collect { a: A ->
            latestB.get()?.let { b -> emit(transform(a, b)) }
        }
    }
}
*/

/**
 * Every time a new value is emitted from A, it will emit instead the latest value from B.
 */
fun <A, B : Any> Flow<A>.switchToLatestFrom(other: Flow<B>): Flow<B> = flow {
    coroutineScope {
        val latestB = atomic<B?>(null)
        val outerScope = this
        launch {
            try {
                other.collect { latestB.set(it) }
            } catch (e: CancellationException) {
                outerScope.cancel(e)
            }
        }
        collect {
            latestB.value?.let { b -> emit(b) }
        }
    }
}

fun <A, B, C : Any?, R> Flow<A>.withLatestFrom(
    other: Flow<B>,
    another: Flow<C>,
    transform: suspend (A, B, C) -> R
): Flow<R> = flow {
    coroutineScope {
        val latestB = atomic<B?>(null)
        val latestC = atomic<C?>(null)
        val outerScope = this
        launch {
            try {
                other.collect { latestB.set(it) }
            } catch (e: CancellationException) {
                outerScope.cancel(e) // cancel outer scope on cancellation exception, too
            }
        }
        launch {
            try {
                another.collect { latestC.set(it) }
            } catch (e: CancellationException) {
                outerScope.cancel(e) // cancel outer scope on cancellation exception, too
            }
        }
        collect { a: A ->
            val b = latestB.value
            val c = latestC.value
            if (b != null && c != null) {
                emit(transform(a, b, c))
            }
        }
    }
}

fun <A, B, C, D : Any?, R> Flow<A>.withLatestFrom(
    first: Flow<B>,
    second: Flow<C>,
    third: Flow<D>,
    transform: suspend (A, B, C, D ) -> R
): Flow<R> = flow {
    coroutineScope {
        val latestB = atomic<B?>(null)
        val latestC = atomic<C?>(null)
        val latestD = atomic<D?>(null)
        val outerScope = this
        launch {
            try {
                first.collect { latestB.set(it) }
            } catch (e: CancellationException) {
                outerScope.cancel(e) // cancel outer scope on cancellation exception, too
            }
        }
        launch {
            try {
                second.collect { latestC.set(it) }
            } catch (e: CancellationException) {
                outerScope.cancel(e) // cancel outer scope on cancellation exception, too
            }
        }
        launch {
            try {
                third.collect { latestD.set(it) }
            } catch (e: CancellationException) {
                outerScope.cancel(e) // cancel outer scope on cancellation exception, too
            }
        }
        collect { a: A ->
            val b = latestB.value
            val c = latestC.value
            val d = latestD.value
            if (b != null && c != null && d != null) {
                emit(transform(a, b, c, d))
            }
        }
    }
}
